# -*- coding: utf-8 -*-
# @Function: 查询测试器
# @Description: 提供问题查询测试功能，根据输入问题进行分词和数据库查询，返回匹配结果
# @Usage: 被all.py导入使用，执行查询测试
# @Dependencies: db_manager.py, tokenizer.py

from D_db.core.db_manager import DBManager
from D_db.core.tokenizer import Tokenizer
from typing import List, Tuple

class QueryTester:
    def __init__(self):
        """初始化查询测试器"""
        self.db_manager = DBManager()
        self.tokenizer = Tokenizer()

    def test_query(self, test_question: str, limit: int = 5):
        """测试查询功能

        Args:
            test_question (str): 测试问题
            limit (int, optional): 返回结果数量限制. Defaults to 5.

        Returns:
            list: 查询结果
        """
        try:
            # 连接数据库
            self.db_manager.connect()

            # 对测试问题进行分词
            token_data = self.tokenizer.tokenize_question(test_question)
            tokens = token_data.get('tokens', [])
            category_id = token_data.get('category_id')  # 获取提取的category_id

            if not tokens:
                print("测试问题分词结果为空")
                return []

            # 查询相关问题，传入category_id参数
            qa_results = self.db_manager.query_qa(tokens, limit, category_id)

            # 输出结果
            print(f"\n测试问题: {test_question}")
            print(f"分词结果: {tokens}")
            if category_id:
                print(f"识别的水库ID: {category_id}")

            # 输出QA查询结果
            print(f"查询到 {len(qa_results)} 个相关问题")
            for i, (question, question_tokens, answer, category_id, context, match_count) in enumerate(qa_results, 1):
                print(f"结果 {i}: {question} (类别: {category_id}, 匹配度: {match_count})")
                print(f"答案: {answer['answer'][:100]}...")
                if context:
                    print(f"上下文: {context}")

            return qa_results

        except Exception as e:
            print(f"查询测试失败: {e}")
            return []
        finally:
            # 关闭数据库连接
            self.db_manager.close()

    def test_queries(self, queries: List[Tuple[str, str]], category_id: str = None):
        """测试多个查询

        Args:
            queries (List[Tuple[str, str]]): 查询列表，每个元组包含一个问题和一个可选的类别ID
            category_id (str, optional): 默认的类别ID. Defaults to None.
        """
        for question, cat_id in queries:
            self.test_query(question, category_id=cat_id or category_id)